package com.yuma.springaichatmodel.controller;

import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@RestController
public class OpenAIChatModelController {

    @Autowired
    private OpenAiChatModel openAiChatModel;

    @GetMapping("/ai/call")
    public String call(@RequestParam(value="msg", defaultValue = "你好") String message){
        return  this.openAiChatModel.call(message);
    }

    // 流式调用
    @GetMapping(value="/ai/stream", produces = "text/html;charset=utf-8")
    public Flux<String> stream(@RequestParam(value="msg", defaultValue = "你好") String message){
        return openAiChatModel.stream(message);
    }

    @GetMapping("/prompt")
    public String prompt(@RequestParam(value="name") String name,@RequestParam("check") String check ,@RequestParam(value="hobby") String hobby){
        // 设置用户输入信息
        UserMessage userMessage = new UserMessage(check);
        // 使用PromptTemplate设置信息
        SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate("你是一个美食咨询助手，可以帮助人们查询美食信息。你得名字是{name}，你应该用你的名字和{hobby}的饮食习惯回复用户的请求。");
        // 替换占位符
        Message systemMessage = systemPromptTemplate.createMessage(Map.of("name", name, "hobby", hobby));
        // 提示词：用户提问和系统身份
        Prompt prompt = new Prompt(List.of(userMessage, systemMessage));
        // 使用chatModel方法
        List<Generation> results = openAiChatModel.call(prompt).getResults();
        return results.stream().map(x->x.getOutput().getText()).collect(Collectors.joining("。"));
    }

    @GetMapping("/prompt2")
    public String prompt2(@RequestParam(value="name") String name,
                          @RequestParam("check") String check ,
                          @RequestParam(value="hobby") String hobby,
                          @Value("classpath:/prompts.st") Resource systemResource){
        // 设置用户输入信息
        UserMessage userMessage = new UserMessage(check);
        // 使用PromptTemplate设置信息
        SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(systemResource);
        // 替换占位符
        Message systemMessage = systemPromptTemplate.createMessage(Map.of("name", name, "hobby", hobby));
        // 提示词：用户提问和系统身份
        Prompt prompt = new Prompt(List.of(userMessage, systemMessage));
        // 使用chatModel方法
        List<Generation> results = openAiChatModel.call(prompt).getResults();
        return results.stream().map(x->x.getOutput().getText()).collect(Collectors.joining(" "));
    }
}
